# Import the necessary libraries
import serial
import os
import sys
import time
import rospy
from sensor_msgs.msg import Imu
from std_msgs.msg import Header
from std_msgs.msg import Int16

# Define a function to convert a four-digit hexadecimal number to a signed integer
def int16_t(num):
    # Check if the integer is greater than or equal to 2^15
    if num >= 2**15:
        # Convert the integer to a signed integer
        signed_int = num - 2**16
    else:
        # The integer is already a signed integer
        signed_int = num
    return signed_int


# Define the serial port and baud rate
ser = serial.Serial(sys.argv[2], 921600)

# Open the serial port
#ser.open()

# Set save path
save_path = os.path.expanduser(sys.argv[1])
# save_path = os.path.expanduser('/media/huarsuma/2FEA-9429/dataset/test')

# Define the file names to save the data
encoder_file_name = 'encoder.txt'
imu_file_name = 'imu.txt'

# Open the files in write mode
with open(os.path.join(save_path, encoder_file_name), 'a') as encoder_file, open(os.path.join(save_path, imu_file_name), 'a') as imu_file:
    while True:
        # Read the data from the serial port
        data = ser.read(1)

        # Check if the data is the start of a frame
        if data == b'\xaa':
            # Read the second byte
            data += ser.read(1)

            # Check if the second byte is the start of a frame
            if data[1] == 0xaf:
                # Read the next 4 bytes
                data += ser.read(3)

                # Check if the frame is encoder data
                if data[2:5] == b'\x01\x01\x02':
                    # Read the next 3 bytes
                    data += ser.read(3)
                    #print(data)

                    # Verify the checksum
                    if sum(data[:-1]) & 0xff == data[-1]:
                        # Extract the encoder data
                        high_byte = data[5]
                        low_byte = data[6]
                        encoder_data = high_byte*256 + low_byte;
                        timestamp = time.time_ns()

                        # Initialize the ROS node
                        try:
                            rospy.init_node('encoder_node', anonymous=True)
                        except rospy.exceptions.ROSException as e:
                            if str(e) == "rospy.init_node() has already been called with different arguments: ('imu_node', ['serial_data.py', '/home/huarsuma/dataset/test'], True, None, False, False)":
                                pass
                            else:
                                raise e

                        # Publish the encoder data to ROS topic
                        pub = rospy.Publisher('encoder_data', Header, queue_size=10)
                        header = Header()
                        header.stamp = rospy.Time.now()
                        header.frame_id = str(encoder_data)
                        pub.publish(header)


                        # Write the data to the encoder file
                        encoder_file.write(f"{timestamp} {encoder_data}\n")
                        
                # Check if the frame is IMU data
                elif data[2:5] == b'\x08\x01\x0C':
                    # Read the next 7 bytes
                    data += ser.read(13)
                    #print(data)

                    # Verify the checksum
                    if sum(data[:-1]) & 0xff == data[-1]:
                        # Extract the IMU data
                        X_acc_low_byte = data[5]
                        X_acc_high_byte = data[6]
                        X_acc = round(int16_t(X_acc_high_byte*256 + X_acc_low_byte)/32768*16*9.8, 10)
                        Y_acc_low_byte = data[7]
                        Y_acc_high_byte = data[8]
                        Y_acc = round(int16_t(Y_acc_high_byte*256 + Y_acc_low_byte)/32768*16*9.8, 10)
                        Z_acc_low_byte = data[9]
                        Z_acc_high_byte = data[10]
                        Z_acc = round(int16_t(Z_acc_high_byte*256 + Z_acc_low_byte)/32768*16*9.8, 10)
                        X_gyro_low_byte = data[11]
                        X_gyro_high_byte = data[12]
                        X_gyro = round(int16_t(X_gyro_high_byte*256 + X_gyro_low_byte)/32768*2000, 10)
                        Y_gyro_low_byte = data[13]
                        Y_gyro_high_byte = data[14]
                        Y_gyro = round(int16_t(Y_gyro_high_byte*256 + Y_gyro_low_byte)/32768*2000, 10)
                        Z_gyro_low_byte = data[15]
                        Z_gyro_high_byte = data[16]
                        Z_gyro = round(int16_t(Z_gyro_high_byte*256 + Z_gyro_low_byte)/32768*2000, 10)
                        timestamp = time.time_ns()

                        # # Initialize the ROS node
                        # try:
                        #     rospy.init_node('imu_node', anonymous=True)
                        # except rospy.exceptions.ROSException as e:
                        #     if str(e) == "rospy.init_node() has already been called with different arguments: ('encoder_node', ['serial_data.py', '/home/huarsuma/dataset/test'], True, None, False, False)":
                        #         pass
                        #     else:
                        #         raise e

                        # # Publish the IMU data to ROS topic
                        # imu_msg = Imu()
                        # imu_msg.header.stamp = rospy.Time.now()
                        # imu_msg.linear_acceleration.x = X_acc
                        # imu_msg.linear_acceleration.y = Y_acc
                        # imu_msg.linear_acceleration.z = Z_acc
                        # imu_msg.angular_velocity.x = X_gyro
                        # imu_msg.angular_velocity.y = Y_gyro
                        # imu_msg.angular_velocity.z = Z_gyro
                        # pub = rospy.Publisher('imu_data', Imu, queue_size=10)
                        # pub.publish(imu_msg)

                        # Write the data to the IMU file
                        imu_file.write(f"{timestamp} {X_gyro} {Y_gyro} {Z_gyro} {X_acc} {Y_acc} {Z_acc}\n")




